# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# Licensed under the 【火山方舟】原型应用软件自用许可协议
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     https://www.volcengine.com/docs/82379/1433703
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import logging

from arkitect.core.component.llm.model import ArkMessage, ArkChatRequest

from deep_search import DeepSearch, ExtraConfig
from search_engine.volc_bot import VolcBotSearchEngine
from search_engine.tavily import TavilySearchEngine
from search_engine.bailian_bot import BailianBotSearchEngine

logging.basicConfig(
    level=logging.INFO, format="[%(asctime)s][%(levelname)s] %(message)s"
)
LOGGER = logging.getLogger(__name__)

# recommend to use DeepSeek-R1 model
REASONING_EP_ID = "deepseek-r1-250120"
# default set to volc bot, if using tavily, change it into "tavily"
SEARCH_ENGINE = "bailian" #"volc_bot"
# optional, if you select tavily as search engine, please configure this
TAVILY_API_KEY = "{YOUR_TAVILY_API_KEY}"
# optional, if you select volc bot as search engine, please configure this
SEARCH_BOT_ID = "bot-20250522214322-zbhw9"
#QUERY = """你是一个资深积木乐高专栏作家，写作流程：基于以下资料，请找出2个和新品发布相关的主题（不要关心品牌，年份和细节），然后搜索相关资料补充（例如，怎样拼装XX,查找积木XX图片等等），写一个吸引人的文章，文章内容包括今天热点和补充的知识资料，务必要保留搜索到的相关图片，不要关注（品牌，定价，用户评价，具体细节），描写热点主题相关能够让观众感到兴趣，不要太认真（资料相关就可以，搜集资料能够写出文章就可以，同时每一个主题至少搜集一次补充资料），现在给你提供小红书的积木今天热点新闻：###  

QUERY = """你是一个资深积木乐高专栏作家，写作流程：
1.基于以下新闻，请找出1个和新闻相关的主题（不要关心品牌，年份等细节，主题应该是简单的5个字以内），然后基于每一个新闻主题整理一个图文介绍（包括展示"cover_url"图片)。例如：新闻"三叉戟的八格车来袭！玛莎拉蒂授权积木实测：咬合力拉满，鲨鱼腮却失踪了？【测评】" ，那么主题就是“八格车”“莎拉蒂”。
2.然后进行至少一次搜索：搜索相关新闻主题进行资料补充（例如，怎样拼装[主题],查找[主题]等等），补充资料也需要图片展示 ，同时整理资料不要复杂（简单补充则可，资料相关就可以，搜集资料能够写出文章就可以，同时每一个主题至少搜集一次补充资料）
3.最后：按以上整理的所有内容写一个吸引人的文章【描写热点主题相关能够让观众感到兴趣和愉快】，文章内容包括今天热点主题介绍+补充的相关知识资料，文章主题结构：【主题1+补充1，主题2+补充2，...】
4.注意：务必要保留搜索到的相关图片，分析主题不要过多关注（品牌，定价，用户评价，具体细节）。
5.现在给你提供今天小红书的积木热点新闻材料：###  

    {
      "_id": "68353ff0ba2d733faeefddd5",
      "note_id": "683261a0000000000303f704",
      "title": "乐高变形金刚G1迷你钢索拼搭图纸分享",
      "user_nickname": "积木星球",
      "user_avatar": "https://sns-avatar-qc.xhscdn.com/avatar/1040g2jo31g5qvkhv1q505pp5k1a083injvulmp8?imageView2/2/w/80/format/jpg",
      "user_id": "6725a0540000000001000e57",
      "liked_count": "303",
      "collected_count": "310",
      "comment_count": "6",
      "shared_count": "15",
      "cover_url": "http://sns-webpic-qc.xhscdn.com/202505271229/f866401723bba6bb8b179d6153126947/1040g2sg31hssj4763ojg5pp5k1a083inlqi7blo!nc_n_webp_mw_1",
      "publish_time": 1748320200,
      "image_arr_new": [
        "https://legoallbricks.oss-cn-shenzhen.aliyuncs.com/xhs_images/1040g2sg31hssj4763ojg5pp5k1a083inlqi7blo!nc_n_webp_prv_1.jpg"
      ],
      "synced": false
    },
    

 ###，"""


async def main():
    search_engine = VolcBotSearchEngine(bot_id=SEARCH_BOT_ID)

    if "tavily" == SEARCH_ENGINE:
        search_engine = TavilySearchEngine(api_key=TAVILY_API_KEY)
    elif "bailian" == SEARCH_ENGINE:
        search_engine = BailianBotSearchEngine(bot_id=SEARCH_BOT_ID)

    deep_research = DeepSearch(
        search_engine=search_engine,
        planning_endpoint_id=REASONING_EP_ID,
        summary_endpoint_id=REASONING_EP_ID,
        extra_config=ExtraConfig(
            max_planning_rounds=5,  # max rounds to planning and search
            max_search_words=5,  # max search words for each planning round.
        )
    )

    thinking = False
    async for chunk in deep_research.astream_deep_research(
            request=ArkChatRequest(model="test",
                                   messages=[ArkMessage(role="user",
                                                        content=QUERY)]),
            question=QUERY
    ):
        if chunk.choices[0].delta.reasoning_content:
            if not thinking:
                print("\n----思考过程----\n")
                thinking = True
            print(chunk.choices[0].delta.reasoning_content, end="")
        elif chunk.choices[0].delta.content:
            if thinking:
                print("\n----输出回答----\n")
                thinking = False
            print(chunk.choices[0].delta.content, end="")


if __name__ == '__main__':
    asyncio.run(main())
